#!/usr/bin/env python3

import subprocess
import time
import itertools

# Grid of parameters to tune
local_loss_weights = [0.1, 1.0]
local_loss_stds = [0.0, 0.1, 0.2]
temps_local_loss = [0.07, 0.15]

# Time interval between job submissions (in seconds)
TIME_INTERVAL = 120
SLURM_SCRIPT = 'multimodal_train.sh'

# Generate all combinations of parameters
param_combinations = list(itertools.product(
    local_loss_weights,
    local_loss_stds,
    temps_local_loss
))

total_jobs = len(param_combinations)
print(f"Total jobs to submit: {total_jobs}")

# Loop over parameter combinations and submit jobs
for idx, params in enumerate(param_combinations):
    local_loss_weight, local_loss_std, temp_local_loss = params

    print(f"Submitting job {idx+1}/{total_jobs} with parameters:")

    # Build the command to submit the SLURM job
    cmd = [
        'sbatch',
        SLURM_SCRIPT,
        str(local_loss_weight),
        str(local_loss_std),
        str(temp_local_loss)
    ]

    # Submit the job
    subprocess.run(cmd)

    # Wait for the specified time interval before submitting the next job
    if idx < total_jobs - 1:
        print(f"Waiting for {TIME_INTERVAL} seconds before submitting the next job...")
        time.sleep(TIME_INTERVAL)
